import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.DERNull;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.pqc.crypto.slhdsa.SLHDSAParameters;
import org.bouncycastle.pqc.crypto.sphincsplus.*;
import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider;
import org.bouncycastle.util.io.pem.PemObject;
import org.bouncycastle.util.io.pem.PemWriter;
import org.bouncycastle.pqc.crypto.slhdsa.*;

import java.io.*;
import java.math.BigInteger;
import java.security.*;
import java.security.cert.X509Certificate;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Date;

public class Main {

    static {
        Security.addProvider(new BouncyCastleProvider());
        Security.addProvider(new BouncyCastlePQCProvider());
    }

    public static void main(String[] args) throws Exception {
        /* 生成SLHDSA共6种张不同强度的 SLH-DSA 证书示例 */
        System.out.println(" /* 生成SLHDSA共6种张不同强度的 SLH-DSA 证书示例 */");
        generateSLHDSACertificate(SLHDSAParameters.sha2_128s,  "sha2_128s");
        generateSLHDSACertificate(SLHDSAParameters.sha2_128f,  "sha2_128f");
        generateSLHDSACertificate(SLHDSAParameters.sha2_192s,  "sha2_192s");
        generateSLHDSACertificate(SLHDSAParameters.sha2_192f,  "sha2_192f");
        generateSLHDSACertificate(SLHDSAParameters.sha2_256s,  "sha2_256s");
        generateSLHDSACertificate(SLHDSAParameters.sha2_256f,  "sha2_256f");
    }

    /**
     * 使用底层 SPHINCS+ API 生成一张自签名 X.509 证书
     * @param params  SLH-DSA 参数集
     * @param suffix  仅用于输出文件名
     */
    private static void generateSLHDSACertificate(SLHDSAParameters params, String suffix) throws Exception {
        SecureRandom  random   = new SecureRandom();
        /* 1. 生成密钥对 */
        SLHDSAKeyPairGenerator kpGen = new SLHDSAKeyPairGenerator();
        kpGen.init(new SLHDSAKeyGenerationParameters(random, params));
        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();
        SLHDSAPublicKeyParameters  pubKey = (SLHDSAPublicKeyParameters)  kp.getPublic();
        SLHDSAPrivateKeyParameters privKey = (SLHDSAPrivateKeyParameters) kp.getPrivate();

        /* 2. 构造 AlgorithmIdentifier（含 NULL） */
        ASN1ObjectIdentifier oid = getSLHDSA_OID(params);
        AlgorithmIdentifier sigAlgId = new AlgorithmIdentifier(oid);//这里oid后面需要跟, DERNull.INSTANCE
        PrivateKeyInfo prikeypkcs8 = new PrivateKeyInfo(sigAlgId,privKey.getEncoded());
        SubjectPublicKeyInfo pubkeypkcs8 = new SubjectPublicKeyInfo(sigAlgId,pubKey.getEncoded());
        /* 3. 构造 SubjectPublicKeyInfo */
        SubjectPublicKeyInfo pubInfo = new SubjectPublicKeyInfo(
                new AlgorithmIdentifier(oid),//这里oid后面需要跟, DERNull.INSTANCE
                pubKey.getEncoded());

        /* 4. 证书元数据 */
        X500Name subject = new X500Name("CN=Test SLH-DSA-" + suffix + " Cert, O=Demo, C=CN");
        BigInteger serial = BigInteger.valueOf(System.currentTimeMillis());
        Instant now = Instant.now();
        Date notBefore = Date.from(now);
        Date notAfter  = Date.from(now.plus(365, ChronoUnit.DAYS));

        X509v3CertificateBuilder certBldr = new X509v3CertificateBuilder(
                subject, serial, notBefore, notAfter, subject, pubInfo);

        /* 5. 自定义 ContentSigner */
        ContentSigner signer = new ContentSigner() {
            private final ByteArrayOutputStream buf = new ByteArrayOutputStream();
            @Override public AlgorithmIdentifier getAlgorithmIdentifier() { return sigAlgId; }
            @Override public OutputStream getOutputStream()             { return buf; }
            @Override public byte[] getSignature() {
                try {
                    SLHDSASigner slhSigner = new SLHDSASigner();
                    slhSigner.init(true, privKey);
                    return slhSigner.generateSignature(buf.toByteArray());
                } catch (Exception e) {
                    throw new RuntimeException("SLH-DSA 签名失败", e);
                }
            }
        };

        /* 6. 构建证书并输出 */
        org.bouncycastle.cert.X509CertificateHolder holder = certBldr.build(signer);
        X509Certificate cert = new JcaX509CertificateConverter()
                .setProvider("BC")
                .getCertificate(holder);

        writeCertificateToPEM(cert,  "./out/slh_dsa_" + suffix + "_cert.crt");
        writePrivateKeyToPEM(prikeypkcs8.getEncoded(), "./out/slh_dsa_" + suffix + "_priv.pem");
        writePublicKeyToPEM (pubkeypkcs8 .getEncoded(),  "./out/slh_dsa_" + suffix + "_pub.pem");

        System.out.println("SLH-DSA-" + suffix + " 证书生成完成");
        System.out.println("  主题 : " + cert.getSubjectX500Principal());
        System.out.println("  签名算法 : " + cert.getSigAlgName());
    }

    /* ------------ 小工具 ------------ */
    private static ASN1ObjectIdentifier getSLHDSA_OID(SLHDSAParameters p) {
        /* 只列常用 6 个，其余同理补全即可 */
        if (p.equals(SLHDSAParameters.sha2_128s)) return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.20");
        if (p.equals(SLHDSAParameters.sha2_128f)) return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.21");
        if (p.equals(SLHDSAParameters.sha2_192s)) return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.22");
        if (p.equals(SLHDSAParameters.sha2_192f)) return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.23");
        if (p.equals(SLHDSAParameters.sha2_256s)) return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.24");
        if (p.equals(SLHDSAParameters.sha2_256f)) return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.25");
        throw new IllegalArgumentException("未定义的 SLH-DSA 参数");
    }

    private static void writeCertificateToPEM(X509Certificate cert, String file) throws Exception {
        try (PemWriter w = new PemWriter(new FileWriter(file))) { w.writeObject(new PemObject("CERTIFICATE", cert.getEncoded())); }
        System.out.println("  证书PEM : " + file);
    }
    private static void writePrivateKeyToPEM(byte[] key, String file) throws Exception {
        try (PemWriter w = new PemWriter(new FileWriter(file))) { w.writeObject(new PemObject("PRIVATE KEY", key)); }
    }
    private static void writePublicKeyToPEM(byte[] key, String file) throws Exception {
        try (PemWriter w = new PemWriter(new FileWriter(file))) { w.writeObject(new PemObject("PUBLIC KEY", key)); }
    }
}